"""
* 名称：model_collision_detection_train_V1.py
* 功能：导入数据集，训练模型，输入36输出1(实际使用模型)
* 注意：
"""

import time

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from tqdm import tqdm

"""********** 1、准备数据集 **********"""
df = pd.read_csv('./dataset/data_set_20240508_cnt.csv', index_col = 0)  # 导入数据
print(df.head())    # 显示数据集前几行，确认数据集格式正确
arr = df.values  # Pandas 对象退化为 NumPy 数组
arr = arr.astype(np.float32)  # 转为 float32 类型数组
ts = torch.tensor(arr)  # 数组转为张量
ts = ts.to('cuda')  # 把训练集搬到 cuda 上
print(f"ts.shape:{ts.shape}")

"""********** 2、划分训练集与测试集 **********"""
train_size = int(len(ts) * 0.7)  # 训练集的样本数量：70%
test_size = len(ts) - train_size  # 测试集的样本数量：30%
ts = ts[torch.randperm(ts.size(0)), :]  # 打乱样本的顺序：非时序
train_Data = ts[: train_size, :]  # 训练集样本
test_Data = ts[train_size:, :]  # 测试集样本

"""********** 3、搭建神经网络模型 **********"""
class DNN(nn.Module):
    """
    * 功能：碰撞检测模型。4个全连接层组成
    * 注意：
    """
    def __init__(self):
        """ 搭建神经网络各层 """
        super(DNN, self).__init__()
        self.net = nn.Sequential(  # 按顺序搭建各层
            nn.Linear(36, 32), nn.Sigmoid(),  # 第 1 层：全连接层
            nn.Linear(32, 8), nn.Sigmoid(),  # 第 2 层：全连接层
            nn.Linear(8, 4), nn.Sigmoid(),  # 第 3 层：全连接层
            nn.Linear(4, 1), nn.Sigmoid()  # 第 4 层：全连接层
        )

    def forward(self, x):
        """ 前向传播 """
        y = self.net(x)  # x 即输入数据
        return y  # y 即输出数据


model = DNN().to('cuda:0')  # 创建子类的实例，并搬到 GPU 上

"""********** 4、损失函数和优化算法选择 **********"""
# 损失函数的选择
loss_fn = nn.BCELoss(reduction = 'mean')

# 优化算法的选择
learning_rate = 0.005  # 设置学习率
optimizer = torch.optim.Adam(model.parameters(), lr = learning_rate)

"""********** 5、训练网络 **********"""
epochs = 800  # 设定训练次数
# cnt = list(range(epochs))
losses = []  # 记录损失函数变化的列表

# 给训练集划分输入与输出
X = train_Data[:, : -1]  # 前 36 列为输入特征
Y = train_Data[:, -1].reshape((-1, 1))  # 后 1 列为输出特征，此处的.reshape((-1,1))将一阶张量升级为二阶张量

start = time.time()
for epoch in tqdm(range(epochs)):
    Pred = model(X)  # 一次前向传播（批量）
    loss = loss_fn(Pred, Y)  # 计算损失函数
    losses.append(loss.item())  # 记录损失函数的变化
    optimizer.zero_grad()  # 清理上一轮滞留的梯度
    loss.backward()  # 一次反向传播
    optimizer.step()  # 优化内部参数
end = time.time()
print("训练耗费{:.2f}s".format(end - start))

Fig = plt.figure()
plt.plot(range(epochs), losses)
plt.ylabel('loss')
plt.xlabel('epoch')
plt.show()

"""********** 6、测试网络 **********"""
# 给测试集划分输入与输出
X = test_Data[:, : -1]  # 前 36 列为输入特征
Y = test_Data[:, -1].reshape((-1, 1))  # 后 1 列为输出特征
with torch.no_grad():  # 该局部关闭梯度计算功能
    Pred = model(X)  # 一次前向传播（批量）
    Pred[Pred >= 0.5] = 1
    Pred[Pred < 0.5] = 0
    correct = torch.sum((Pred == Y).all(1))  # 预测正确的样本
    total = Y.size(0)  # 全部的样本数量
    print(f'测试集精准度: {100 * correct / total} %')

"""********** 7、保存模型 **********"""
# 保存网络
torch.save(model, 'model_collision_detection_v1.pth')





# 把模型赋给新网络
# new_model = torch.load('DNN_model.pth')
# with torch.no_grad(): # 该局部关闭梯度计算功能
#     Pred = new_model(X) # 一次前向传播（批量）
#     Pred[Pred>=0.5] = 1
#     Pred[Pred<0.5] = 0


# X = [1.55, 0, 0, 0, 0, 0, 0.040, -0.49, 1.16, 0.028983988, -0.272166371, 1.22, -0.013745327, 0.071199894, 1.241997838,
#      -0.107642092, -0.195327416, 1.293527722, -0.220300868, -0.3651793, 1.408981323, -0.341132879, -0.499009162,
#      1.523368001, 0.149345234, -0.141609564, 1.127255917, 0.274599493, -0.312070817, 1.064984202, 0.323910713,
#      -0.395466715, 0.906407595, 0.021285573, -0.113410875, 1.236373544]
# # X = X.astype(np.float32)
# ts = torch.tensor(X)
# ts = ts.to('cuda')
# print(ts.shape)
#
# with torch.no_grad():  # 该局部关闭梯度计算功能
#     Pred = model(ts)  # 一次前向传播（批量）
#     Pred[Pred >= 0.5] = 1
#     Pred[Pred < 0.5] = 0
#     print(Pred)

# 训练网络
epochs = 100
losses = [] # 记录损失函数变化的列表
X = torch.rand(256,1,72)
X = torch.tensor(X)
X = X.to('cuda:0')
print(f"X.shape:{X.shape}")
for epoch in range(epochs):
    Pred = model(X)  # 一次前向传播（批量）
    loss = loss_fn(Pred, Y)  # 计算损失函数
    losses.append(loss.item())  # 记录损失函数的变化
    optimizer.zero_grad()  # 清理上一轮滞留的梯度
    loss.backward()  # 一次反向传播
    optimizer.step()  # 优化内部参数
Fig = plt.figure()
plt.plot(range(len(losses)), losses)
plt.show()

